import random as rand
from BinomialEdgeGen import BinomialEdgeGen

class ERClustersEdgeGen(object):

    def __init__(self, maxClusterSize, edgeProb):
        '''
        Constructor
        '''
        
        self.MaxClusterSize = maxClusterSize
        self.EdgeProb = edgeProb
        
        
    def Apply(self, network):
        ERGenerator = BinomialEdgeGen(self.EdgeProb)
        
        ERGenerator.Apply(network)
        
        #Make the edges bidirectional
        for edge in network.graph.edges_iter():
            network.graph.add_edge(edge[1], edge[0])

        curNodeIndex = 0
        numNodes = len(network.agentList)
        while(curNodeIndex < numNodes):
            curClusterSize = rand.randint(1, self.MaxClusterSize)
            
            if (curNodeIndex + curClusterSize) > numNodes:
                curClusterSize = numNodes - curNodeIndex
            
            nodeAtStartOfNextClusterIndex = curNodeIndex + curClusterSize
            for nodeUIndex in range(curNodeIndex, nodeAtStartOfNextClusterIndex):
                nodeU = network.agentList[nodeUIndex]
                for nodeVIndex in range(curNodeIndex, nodeAtStartOfNextClusterIndex):
                    if nodeUIndex != nodeVIndex:
                        
                        nodeV = network.agentList[nodeVIndex]
                        
                        network.graph.add_edge(nodeU, nodeV)
                        network.graph.add_edge(nodeV, nodeU)

            curNodeIndex += curClusterSize
        
            
        
    